/**************************************************************************************************
* Title: SumHeapAllocator.cpp
* Author: Gael Huber
* Description: General purpose heap allocator
**************************************************************************************************/
#include "SumHeapAllocator.h"

/**************************************************************************************************
* Allocate a block of memory
**************************************************************************************************/
void* const HeapAllocator::alloc(size_t bytes) {
	// First make sure the requested size is at least the smallest quantity
	if(bytes < minSize)
		bytes = 32;
	
	// Then, find the smallest compatible node
	LengthTree::Node* ln = lengthTree->findSmallestContaining(bytes);

	// Check that it isn't 0
	assert(ln != 0);

	// Find the node representative of address
	PointerTree::Node* pn = pointerTree->find((MemAddress) ln - 16);

	// Address to return
	void* toReturn = (void*) pn->getAddress();

	// Split the memory chunk
	splitBlock(bytes, pn, ln);

	// Return the address
	return toReturn;
}

/**************************************************************************************************
* Free a block of memory
**************************************************************************************************/
void HeapAllocator::free(void* address) {
	// Get the length of the block we are freeing
	size_t* length = (size_t*) ((char*) address - 4);

	// Find the address of the next block of memory
	MemAddress nextBlock = (MemAddress)address + *length + 4;

	// See whether that block is free
	PointerTree::Node* pn = pointerTree->find(nextBlock);
	if(pn != 0) {
		// The succeeding block is free, so we should merge this block with that one
		// The first step is to find the corresponding length block
		LengthTree::Node* ln = (LengthTree::Node*) ((char*) pn + 16);

		// Remove both nodes from their respective trees
		pointerTree->remove(pn);
		lengthTree->remove(ln);

		// Merge the two blocks
		mergeBlocks(address, pn);

	} else {
		// It doesn't exist, so we must create a new free block and insert it into both trees
		PointerTree::Node* pn = (PointerTree::Node*) ((char*) address);
		*pn = PointerTree::Node((MemAddress) address);
		pointerTree->insert(pn);

		LengthTree::Node* ln = (LengthTree::Node*) ((char*) address + 16);
		*ln = LengthTree::Node(*length);
		lengthTree->insert(ln);
	}
}

/**************************************************************************************************
* Initialize the heap allocator
**************************************************************************************************/
void HeapAllocator::initialize(void) {
	// Allocate the total size
	blockStart = malloc(totalAllocated);

	// Designate space for the preblock and specify size of block
	size_t* blockSize = (size_t*) ((char*) blockStart);
	*blockSize = totalAllocated - 4;

	// Initialize the two trees
	pointerTree = new PointerTree();
	lengthTree = new LengthTree();

	// Create the root nodes
	PointerTree::Node* pn = (PointerTree::Node*) ((char*) blockStart + 4);	// Initial block starts 4 
		// after the start (pre-block data containing size of block)
	*pn = PointerTree::Node((MemAddress) blockStart + 4);
	pointerTree->insert(pn);

	LengthTree::Node* ln = (LengthTree::Node*) ((char*) blockStart + 20);	// Initial block starts 20
		// after the start (pre-block data containing the size of block)
	*ln = LengthTree::Node(totalAllocated - 4);
	lengthTree->insert(ln);
}

/**************************************************************************************************
* Split a block of memory of size bytes and the remainder
**************************************************************************************************/
void HeapAllocator::splitBlock(size_t bytes, PointerTree::Node* pn, LengthTree::Node* ln) {
	// Remove each node from their respective trees
	pointerTree->remove(pn);
	lengthTree->remove(ln);

	// Split the memory chunk:
	// Move forward by n bytes and designate this as the end of the block. The size of the new block
	// will be length - n bytes. Place this value in the new memory space and create two new nodes out
	// of the succeeding space. Insert these new nodes into their respective trees.

	// Get the old memory position and update its length
	size_t* length = (size_t*) ((char*) pn - 4);
	*length = bytes;

	// Get the new memory position
	char* position = (char*) pn + bytes;

	// Set the new length at that position
	length = (size_t*) position;
	*length = ln->getLength() - bytes - 4;

	// Designate a new pointer node and length node
	PointerTree::Node* npn = (PointerTree::Node*) (position + 4);
	*npn = PointerTree::Node((MemAddress) position + 4);
	pointerTree->insert(npn);

	LengthTree::Node* nln = (LengthTree::Node*) ((char*) position + 20);
	*nln = LengthTree::Node(*length);
	lengthTree->insert(nln);
}

/**************************************************************************************************
* Merge two blocks at the corresponding addresses
**************************************************************************************************/
void HeapAllocator::mergeBlocks(void* a, void* b) {
	// Find the length of the first and second blocks
	size_t* length1 = (size_t*) ((char*) a - 4);
	size_t* length2 = (size_t*) ((char*) b - 4);

	// Set the new size of the block, the combination of the two blocks plus the additional 4 bytes 
	// gained from merging due to one less storage requirement (only store size for the block once)
	*length1 = *length1 + *length2 + 4;

	// Create the new nodes
	PointerTree::Node* pn = (PointerTree::Node*) a;
	*pn = PointerTree::Node((MemAddress)a);
	pointerTree->insert(pn);

	LengthTree::Node* ln = (LengthTree::Node*) ((char*) a + 16);
	*ln = LengthTree::Node(*length1);
	lengthTree->insert(ln);
}